package com.woniu.ai;

import okhttp3.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.Map;
import java.util.Scanner;

public class StreamAi {
    public static void doPost(String url, String json) {
        MediaType mediaType = MediaType.parse("application/json; charset=utf-8");
        Request requestbuilder = new Request.Builder()
                .url(url)
                .post(RequestBody.create(mediaType, json))
                .build();

        OkHttpClient okHttpClient = new OkHttpClient();

        // 针对异步响应进行处理
        okHttpClient.newCall(requestbuilder).enqueue(new Callback() {
            @Override
            public void onResponse(Call call, Response response) throws IOException {
                try (ResponseBody responseBody = response.body()) {
                    if (responseBody != null) {
                        processStream(responseBody);   // 调用processStream处理响应流
                    }
                }
            }
            @Override
            public void onFailure(Call call, IOException e) {
                e.printStackTrace();
            }
        });
    }
    public static void processStream(ResponseBody responseBody) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        okio.BufferedSource source = responseBody.source();
        while (!source.exhausted()) {
            String line = source.readUtf8Line();
            if (line != null) {
                Map<String, String> map = mapper.readValue(line, Map.class);
                System.out.print(map.get("response"));  // 此处的输入不能换行
            }
        }
        System.exit(0);
    }
    public static void main(String[] args) {
            String url = "http://localhost:11434/api/generate";
            System.out.println("请输入问题");
            Scanner scanner = new Scanner(System.in);
            String message = scanner.nextLine();
            String json = "{'model': 'deepseek-r1:1.5b','prompt': '%s', 'stream': true}"
                    .formatted(message).replace("'", "\"");
            new StreamAi().doPost(url, json);
    }
}
